package Test;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.IOException;
import java.net.URL;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;


/**
 * User: RuanYongliu
 * Date: 15-1-19
 * Time: 下午7:56
 */
public class ImageHelper {

    // 实现高斯模糊开了多线程，如果不用线程池有可能造成资源被消耗光的情况
    private static ExecutorService gaussianThreadPool = new ThreadPoolExecutor(
            AdCommonConst.GaussianBlurConst.CORE_POOL_SIZE,
            AdCommonConst.GaussianBlurConst.MAX_POOL_SIZE,
            0, TimeUnit.SECONDS,
            new ArrayBlockingQueue<Runnable>(AdCommonConst.GaussianBlurConst.WORK_QUEUE_SIZE, true),
            new ThreadPoolExecutor.AbortPolicy());

    private static class GaussianBlurColThread extends Thread {
        int n, r, kernelRadius, width, height, st, ed;
        int[] rawA, rawR, rawG, rawB;
        double[] middleA, middleR, middleG, middleB;
        double[][] kernel;
        CountDownLatch latch;
        public GaussianBlurColThread(int r, int kernelRadius, int width, int height, int st, int ed,
                                     double[][] kernel, double[] middleA, double[] middleR, double[] middleG,
                                     double[] middleB, int[] rawA, int[] rawR, int[] rawG, int[] rawB,
                                     CountDownLatch latch) {
            this.n = rawA.length;
            this.r = r;
            this.kernelRadius = kernelRadius;
            this.width = width;
            this.height = height;
            this.st = st;
            this.ed = ed;
            this.kernel = kernel;
            this.rawA = rawA;
            this.rawR = rawR;
            this.rawG = rawG;
            this.rawB = rawB;
            this.middleA = middleA;
            this.middleR = middleR;
            this.middleG = middleG;
            this.middleB = middleB;
            this.latch = latch;
        }

        @Override
        public void run() {
            for (int i = st; i <= ed; ++i) {
                int x = i / width, y = i % width;
                middleA[i] = middleR[i] = middleG[i] = middleB[i] = 0;
                double s = 0;
                for (int offsetX = -r; offsetX <= r; ++offsetX) {
                    int rx = Math.max(0, Math.min(height - 1, x + offsetX));
                    int j = rx * width + y;
                    middleA[i] += rawA[j] * kernel[offsetX + kernelRadius][kernelRadius];
                    middleR[i] += rawR[j] * kernel[offsetX + kernelRadius][kernelRadius];
                    middleG[i] += rawG[j] * kernel[offsetX + kernelRadius][kernelRadius];
                    middleB[i] += rawB[j] * kernel[offsetX + kernelRadius][kernelRadius];
                    s += kernel[offsetX + kernelRadius][kernelRadius];
                }
                middleA[i] /= s;
                middleR[i] /= s;
                middleG[i] /= s;
                middleB[i] /= s;
            }
            latch.countDown();
        }
    }

    private static class GaussianBlurRowThread extends Thread {
        int n, r, kernelRadius, width, height, st, ed;
        int[] blurred;
        double[] middleA, middleR, middleG, middleB;
        double[][] kernel;
        CountDownLatch latch;

        public GaussianBlurRowThread(int r, int kernelRadius, int width, int height, int st, int ed,
                                     double[][] kernel, double[] middleA, double[] middleR, double[] middleG,
                                     double[] middleB, int[] blurred, CountDownLatch latch) {
            this.n = blurred.length;
            this.r = r;
            this.kernelRadius = kernelRadius;
            this.width = width;
            this.height = height;
            this.st = st;
            this.ed = ed;
            this.blurred = blurred;
            this.kernel = kernel;
            this.middleA = middleA;
            this.middleR = middleR;
            this.middleG = middleG;
            this.middleB = middleB;
            this.latch = latch;
        }

        @Override
        public void run() {
            double[] blurredA = new double[n], blurredR = new double[n], blurredG = new double[n],
                    blurredB = new double[n];
            for (int i = st; i <= ed; ++i) {
                int x = i / width, y = i % width;
                blurredA[i] = blurredR[i] = blurredG[i] = blurredB[i] = 0;
                double s = 0;
                for (int offsetY = -r; offsetY <= r; ++offsetY) {
                    int ry = Math.max(0, Math.min(width - 1, y + offsetY));
                    int j = x * width + ry;
                    blurredA[i] += middleA[j] * kernel[kernelRadius][offsetY + kernelRadius];
                    blurredR[i] += middleR[j] * kernel[kernelRadius][offsetY + kernelRadius];
                    blurredG[i] += middleG[j] * kernel[kernelRadius][offsetY + kernelRadius];
                    blurredB[i] += middleB[j] * kernel[kernelRadius][offsetY + kernelRadius];
                    s += kernel[kernelRadius][kernelRadius + offsetY];
                }
                blurredA[i] /= s;
                blurredR[i] /= s;
                blurredG[i] /= s;
                blurredB[i] /= s;
            }

            for (int i = st; i < ed; ++i) {
                blurred[i] = getARGB(blurredA[i], blurredR[i], blurredG[i], blurredB[i]);
            }
            latch.countDown();
        }
    }

    private static int getARGB(double alpha, double red, double green, double blue) {
        return (clamp(alpha) << 24) + (clamp(red) << 16) + (clamp(green) << 8) + clamp(blue);
    }

    private static int clamp(int argb) {
        return argb & 0xFF;
    }

    private static int clamp(double argb) {
        return clamp((int)argb);
    }

    /**
     * 实现高斯模糊
     * @param image
     */
    public static BufferedImage gaussianBlur(BufferedImage image, final int r, final int kernelRadius,
                                             final double[][] kernel, int threadNum)
            throws IllegalArgumentException, IOException, InterruptedException {

        final int width = image.getWidth();
        final int height = image.getHeight();
        // 读取像素
        int type= image.getType();
        int[] pixels;
        if (type == BufferedImage.TYPE_INT_ARGB || type == BufferedImage.TYPE_INT_RGB) {
            pixels = (int [])image.getRaster().getDataElements(0, 0, width, height, null);
        } else {
            pixels = image.getRGB(0, 0, width, height, null, 0, width);
        }
        if (pixels == null) {
            throw new IllegalArgumentException("bad image");
        }

        // 模糊过程, 并行计算
        // 算法: 先按列实现一次一维高斯模糊，再按行实现一次一维高斯模糊，复杂度h*w*(2*r+1)*2
        // 优化，并行运算
        if (threadNum < 0) {
            threadNum = 1;
        }

        int n = pixels.length, threadSize = n / threadNum;
        double[] middleA = new double[n], middleR = new double[n], middleG = new double[n],
                middleB = new double[n];
        int[] rawA = new int[n], rawR = new int[n], rawG = new int[n], rawB = new int[n];
        for (int i = 0; i < n; ++i) {
            rawA[i] = clamp(pixels[i] >> 24);
            rawR[i] = clamp(pixels[i] >> 16);
            rawG[i] = clamp(pixels[i] >> 8);
            rawB[i] = clamp(pixels[i]);
        }
        Thread[] threads = new Thread[threadNum];
        CountDownLatch latch = new CountDownLatch(threadNum);
        for (int i = 0; i < threadNum; ++i) {
            int st = i * threadSize;
            int ed = i + 1 == threadNum? n - 1: (i + 1) * threadSize - 1;
            threads[i] = new GaussianBlurColThread(r, kernelRadius, width, height, st, ed, kernel, middleA, middleR,
                    middleG, middleB, rawA, rawR, rawG, rawB, latch);
            gaussianThreadPool.execute(threads[i]);
        }
        latch.await();

        latch = new CountDownLatch(threadNum);
        int[] blurred = new int[n];
        for (int i = 0; i < threadNum; ++i) {
            int st = i * threadSize;
            int ed = i + 1 == threadNum? n - 1: (i + 1) * threadSize - 1;
            threads[i] = new GaussianBlurRowThread(r, kernelRadius, width, height, st, ed, kernel, middleA,
                    middleR, middleG, middleB, blurred, latch);
            gaussianThreadPool.execute(threads[i]);
        }
        latch.await();

        // 写成新图片
        BufferedImage blurredImage = new BufferedImage(width, height, type);
        if (type == BufferedImage.TYPE_INT_ARGB || type == BufferedImage.TYPE_INT_RGB) {
            blurredImage.getRaster().setDataElements(0, 0, width, height, blurred);
        } else {
            blurredImage.setRGB(0, 0, width, height, blurred, 0, width);
        }
        return blurredImage;
    }

    public static BufferedImage getImage(String url) {
        try {
            BufferedImage image = ImageIO.read(new URL(url));
            return image;
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;
    }
}
